{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import os\n",
    "import glob\n",
    "%matplotlib inline\n",
    "from scipy.stats import zscore\n",
    "import imp\n",
    "# modules\n",
    "import decoders, utils, tuning, learning, mainfigs, suppfigs\n",
    "\n",
    "### WHERE YOU DOWNLOADED THE FIGSHARE\n",
    "dataroot = '/home/carsen/dm11/releases/stringer-et-al-2019a/fs125/'\n",
    "\n",
    "# file list\n",
    "db = np.load(os.path.join(dataroot, 'database.npy'), allow_pickle=True)\n",
    "fs = []\n",
    "for di in db:\n",
    "    mname = di['mouse_name']\n",
    "    datexp = di['date']\n",
    "    blk = di['block']\n",
    "    stype = di['expt']\n",
    "    \n",
    "    fname = '%s_%s_%s_%s.npy'%(stype, mname, datexp, blk)\n",
    "    fs.append(os.path.join(dataroot, fname))\n",
    "\n",
    "### WHERE YOU WANT TO SAVE THE OUTPUTS OF THE ANALYSIS SCRIPTS AND THE FIGURES (if save_figure=True)\n",
    "saveroot = '/media/carsen/DATA1/2P/ori_data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#~~~~~~~~~~~~~~~ DATA ANALYSIS ~~~~~~~~~~~~~~~~~~~~~~~~~#\n",
    "\n",
    "\n",
    "################### DECODING ANGLES 0-360 #######################\n",
    "dstr = ['independent', 'linear']\n",
    "for dtype in range(2):\n",
    "    print('>>>>>>>>>> %s <<<<<<<<<<<<<<'%dstr[dtype])\n",
    "    if dtype==0:\n",
    "        islinear=False\n",
    "    else:\n",
    "        islinear=True\n",
    "        \n",
    "    ### DECODER\n",
    "    E, errors, stims, SNR, theta_pref = decoders.run_decoder(fs[:27], linear=islinear)\n",
    "    if dtype==1:\n",
    "        np.save(os.path.join(saveroot,'%s_decoder_all_stims.npy'%dstr[dtype]), \n",
    "                {'E': E, 'errors': errors, 'stims': stims})\n",
    "    else:\n",
    "        np.save(os.path.join(saveroot,'%s_decoder_all_stims.npy'%dstr[dtype]), \n",
    "                {'E': E, 'errors': errors, 'stims': stims, 'SNR': SNR, 'theta_pref': theta_pref})\n",
    "\n",
    "    ### ASYMPTOTIC ANALYSIS\n",
    "    E, ccE, nsplit, npop, nstim, E2 = decoders.asymptotics(fs[:6], linear=islinear)\n",
    "    np.save(os.path.join(saveroot,'%s_decoder_asymp.npy'%dstr[dtype]), \n",
    "                {'E': E, 'ccE': ccE, 'nsplit': nsplit, 'npop': npop, 'nstim': nstim, 'E2': E2})\n",
    "\n",
    "### decoding without subtracting spont\n",
    "E, errors, stims, SNR, theta_pref = decoders.run_decoder(fs[:27], linear=islinear, npc=0)\n",
    "np.save(os.path.join(saveroot,'%s_decoder_with_spont.npy'%dstr[dtype]), \n",
    "                {'E': E, 'errors': errors, 'stims': stims})\n",
    "\n",
    "### linear decoding from PC's\n",
    "nPC = 2**np.arange(2,13,1)\n",
    "errors, apreds, atrues = decoders.pc_decoding(fs[:6], nPC)\n",
    "np.save(os.path.join(saveroot, 'pcdecode.npy'), \n",
    "        {'errors': errors, 'apreds': apreds, 'atrues': atrues, 'nPC': nPC})\n",
    " \n",
    "\n",
    "#################### SINGLE NEURON STATS ########################\n",
    "\n",
    "# load a dataset for stim distances\n",
    "dat = np.load(fs[0], allow_pickle=True).item()\n",
    "sresp, istim, itrain, itest = utils.compile_resp(dat, npc=npc)\n",
    "\n",
    "### STIMULUS DISTANCES\n",
    "cc, dtheta_aligned, cbinned, embedding = tuning.population_distances(sresp, istim)\n",
    "np.save(os.path.join(saveroot,'popdist.npy'), \n",
    "        {'cc':cc, 'dtheta_aligned': dtheta_aligned, 'cbinned':cbinned,'embedding':embedding, 'istim': istim})\n",
    "\n",
    "###### population tuning curves\n",
    "d = np.load(os.path.join(saveroot, 'independent_decoder_all_stims.npy'), allow_pickle=True).item()\n",
    "angle_pref = d['theta_pref']\n",
    "### static gratings population tuning curves\n",
    "avg_tuning, tbins = tuning.population_tuning(fs[:6], angle_pref[:6], saveroot)\n",
    "np.save(os.path.join(saveroot,'avgneur_static.npy'), \n",
    "        {'avg_tuning': avg_tuning, 'thetas': angle_pref[:6], 'tbins': tbins})\n",
    "\n",
    "### drifting gratings population tuning curves\n",
    "avg_tuning, tbins = tuning.population_tuning(fs[18:21], angle_pref[18:21], saveroot)\n",
    "np.save(os.path.join(saveroot,'avgneur_drifting.npy'), \n",
    "        {'avg_tuning': avg_tuning, 'thetas': angle_pref[18:21], 'tbins': tbins})\n",
    "    \n",
    "\n",
    "    \n",
    "\n",
    "################ DISCRIMINATION TASK #######################\n",
    "print('DISCRIMINATION')\n",
    "\n",
    "### linear discriminator (all datasets except biased)\n",
    "P,d75,drange = decoders.run_discrimination(fs[:27], decoder='linear')\n",
    "np.save(os.path.join(saveroot, 'linear_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "### linear discriminator on biased datasets (43-47 degree stimuli shown only)\n",
    "npop,nstim,Pstim,Pneur,drange = decoders.dense_discrimination(fs[27:])\n",
    "np.save(os.path.join(saveroot, 'dense_discrimination.npy'), \n",
    "        {'npop': npop, 'nstim': nstim, 'Pneur': Pneur, 'Pstim': Pstim, 'drange': drange})\n",
    "\n",
    "###  discrimination in passive versus running trials\n",
    "all_running = np.load(os.path.join(dataroot, 'all_running.npy'))\n",
    "P,d75,drange = decoders.runspeed_discrimination(fs[:6], all_running[:6])\n",
    "np.save(os.path.join(saveroot, 'runspeed_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "### depth discrimination (split into layers 2/3 and layer 4)\n",
    "all_depths = np.load(os.path.join(dataroot, 'all_depths.npy'))\n",
    "P,d75,drange = decoders.layer_discrimination(fs[:6], all_depths[:6])\n",
    "np.save(os.path.join(saveroot, 'layer_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "### chron discrimination (split train/test in time)\n",
    "P,d75,drange = decoders.chron_discrimination(fs[:6], all_depths[:6])\n",
    "np.save(os.path.join(saveroot, 'chron_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "### neural network discriminator\n",
    "P,d75,drange = decoders.run_discrimination(fs[:6], decoder='deep_net')\n",
    "np.save(os.path.join(saveroot, 'nn_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "### random forest discriminator\n",
    "P,d75,drange = decoders.run_discrimination(fs[:6], decoder='random_forest')\n",
    "np.save(os.path.join(saveroot, 'rf_discrimination.npy'), {'P': P, 'd75': d75, 'drange': drange})\n",
    "\n",
    "#~~~~~~~~~~~~~~~ PERCEPTRONS ~~~~~~~~~~~~~~~~~~~~~~~~~#\n",
    "\n",
    "### PERCEPTRONS\n",
    "# train on easy task\n",
    "nstim32, perf32 = learning.train_perceptrons(fs[:6], task_type='easy')\n",
    "# train on hard task\n",
    "nstim, perf = learning.train_perceptrons(fs[27:], task_type='hard')\n",
    "np.save(os.path.join(saveroot, 'strong_learn.npy'), \n",
    "         {'nstim32': nstim32, 'perf32': perf32, 'perf': perf, 'nstim': nstim})\n",
    "\n",
    "\n",
    "### WEAK LEARNERS\n",
    "# train on easy task\n",
    "P, drange, ccN = learning.train_weak_learners(fs[:6])\n",
    "np.save(os.path.join(saveroot, 'weak_learn.npy'), {'P': P, 'drange': drange, 'ccN': ccN})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%config InlineBackend.close_figures = False\n",
    "%config InlineBackend.print_figure_kwargs={'bbox_inches':None}\n",
    "\n",
    "imp.reload(mainfigs)\n",
    "imp.reload(suppfigs)\n",
    "\n",
    "#### MAIN FIGURES\n",
    "\n",
    "#### SUPPLEMENTARY FIGURES\n",
    "# if you want to save the figures, add an optional input save_figure=True\n",
    "\n",
    "f = mainfigs.fig1(dataroot, saveroot)\n",
    "\n",
    "f = mainfigs.fig2(dataroot, saveroot)\n",
    "\n",
    "f = mainfigs.fig3(dataroot,saveroot)\n",
    "\n",
    "f = mainfigs.fig4(dataroot, saveroot)\n",
    "\n",
    "f = suppfigs.population_curves(saveroot)\n",
    "\n",
    "f = suppfigs.stim_distances(saveroot)\n",
    "\n",
    "f = suppfigs.pc_errors(saveroot)\n",
    "\n",
    "f = suppfigs.linear_decoders(dataroot, saveroot)\n",
    "\n",
    "f = suppfigs.asymptotics(saveroot)\n",
    "\n",
    "f = suppfigs.spont_sub(saveroot)\n",
    "\n",
    "f = suppfigs.discr_chron(saveroot)\n",
    "\n",
    "f = suppfigs.discr_layers(saveroot)\n",
    "\n",
    "f = suppfigs.discr_nn_rf(saveroot)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
